package ast

import (
	"fmt"
	"strings"
)

// Node represents an AST node
type Node interface {
	String() string
}

// Expression represents any Scheme expression
type Expression interface {
	Node
	expressionNode()
}

// Program represents the root of the AST
type Program struct {
	Expressions []Expression
}

func (p *Program) String() string {
	var out strings.Builder
	for _, expr := range p.Expressions {
		out.WriteString(expr.String())
		out.WriteString("\n")
	}
	return out.String()
}

// Symbol represents a symbol/identifier
type Symbol struct {
	Value string
}

func (s *Symbol) String() string {
	return s.Value
}

func (s *Symbol) expressionNode() {}

// StringLiteral represents a string literal
type StringLiteral struct {
	Value string
}

func (sl *StringLiteral) String() string {
	return fmt.Sprintf("\"%s\"", sl.Value)
}

func (sl *StringLiteral) expressionNode() {}

// IntegerLiteral represents an integer literal
type IntegerLiteral struct {
	Value int64
}

func (il *IntegerLiteral) String() string {
	return fmt.Sprintf("%d", il.Value)
}

func (il *IntegerLiteral) expressionNode() {}

// RealLiteral represents a real (floating-point) literal
type RealLiteral struct {
	Value float64
}

func (rl *RealLiteral) String() string {
	return fmt.Sprintf("%g", rl.Value)
}

func (rl *RealLiteral) expressionNode() {}

// BooleanLiteral represents a boolean literal
type BooleanLiteral struct {
	Value bool
}

func (bl *BooleanLiteral) String() string {
	if bl.Value {
		return "#t"
	} else {
		return "#f"
	}
}

func (bl *BooleanLiteral) expressionNode() {}

// List represents a list of expressions
type List struct {
	Elements []Expression
}

func (l *List) String() string {
	var out strings.Builder
	out.WriteString("(")
	for i, elem := range l.Elements {
		if i > 0 {
			out.WriteString(" ")
		}
		out.WriteString(elem.String())
	}
	out.WriteString(")")
	return out.String()
}

func (l *List) expressionNode() {}

// DefineExpr represents a define special form
type DefineExpr struct {
	Name  Symbol
	Value Expression
}

func (d *DefineExpr) String() string {
	return fmt.Sprintf("(define %s %s)", d.Name.String(), d.Value.String())
}

func (d *DefineExpr) expressionNode() {}

// DisplayExpr represents a display special form
type DisplayExpr struct {
	Value Expression
}

func (de *DisplayExpr) String() string {
	return fmt.Sprintf("(display %s)", de.Value.String())
}

func (de *DisplayExpr) expressionNode() {}

// ArithmeticExpr represents an arithmetic expression
type ArithmeticExpr struct {
	Operator string
	Operands []Expression
}

func (ae *ArithmeticExpr) String() string {
	var out strings.Builder
	out.WriteString("(")
	out.WriteString(ae.Operator)
	for _, operand := range ae.Operands {
		out.WriteString(" ")
		out.WriteString(operand.String())
	}
	out.WriteString(")")
	return out.String()
}

func (ae *ArithmeticExpr) expressionNode() {}

// LambdaExpr represents a lambda expression
type LambdaExpr struct {
	Parameters []Symbol
	Body       []Expression
}

func (l *LambdaExpr) String() string {
	var out strings.Builder
	out.WriteString("(lambda (")
	for i, param := range l.Parameters {
		if i > 0 {
			out.WriteString(" ")
		}
		out.WriteString(param.String())
	}
	out.WriteString(") ")
	for i, expr := range l.Body {
		if i > 0 {
			out.WriteString(" ")
		}
		out.WriteString(expr.String())
	}
	out.WriteString(")")
	return out.String()
}

func (l *LambdaExpr) expressionNode() {}

// FunctionCall represents a function call
type FunctionCall struct {
	Function  Expression
	Arguments []Expression
}

func (fc *FunctionCall) String() string {
	var out strings.Builder
	out.WriteString("(")
	out.WriteString(fc.Function.String())
	for _, arg := range fc.Arguments {
		out.WriteString(" ")
		out.WriteString(arg.String())
	}
	out.WriteString(")")
	return out.String()
}

func (fc *FunctionCall) expressionNode() {}

// BeginExpr represents a begin special form
// (begin expr1 expr2 ... exprN) evaluates each expression in order
// and returns the value of the last expression
// according to R4RS specification

// BeginExpr represents a begin special form
type BeginExpr struct {
	Body []Expression
}

func (b *BeginExpr) String() string {
	var out strings.Builder
	out.WriteString("(begin")
	for _, expr := range b.Body {
		out.WriteString(" ")
		out.WriteString(expr.String())
	}
	out.WriteString(")")
	return out.String()
}

func (b *BeginExpr) expressionNode() {}

// SetExpr represents a set! special form
// (set! var value) modifies the binding of an existing variable
// according to R4RS specification
type SetExpr struct {
	Name  Symbol
	Value Expression
}

func (s *SetExpr) String() string {
	return fmt.Sprintf("(set! %s %s)", s.Name.String(), s.Value.String())
}

func (s *SetExpr) expressionNode() {}

// LetExpr represents a let special form
// (let ((var1 val1) (var2 val2) ...) body...) creates local bindings
// according to R4RS specification
type LetExpr struct {
	Bindings []LetBinding // Variable bindings
	Body     []Expression // Body expressions
}

// LetBinding represents a single binding in a let expression
type LetBinding struct {
	Name  Symbol
	Value Expression
}

func (l *LetExpr) String() string {
	var out strings.Builder
	out.WriteString("(let (")
	for i, binding := range l.Bindings {
		if i > 0 {
			out.WriteString(" ")
		}
		out.WriteString("(")
		out.WriteString(binding.Name.String())
		out.WriteString(" ")
		out.WriteString(binding.Value.String())
		out.WriteString(")")
	}
	out.WriteString(") ")
	for i, expr := range l.Body {
		if i > 0 {
			out.WriteString(" ")
		}
		out.WriteString(expr.String())
	}
	out.WriteString(")")
	return out.String()
}

func (l *LetExpr) expressionNode() {}

// IfExpr represents an if special form
// (if condition consequence alternative) evaluates condition and returns consequence if true, alternative if false
// according to R4RS specification
// If alternative is not provided, nil is returned if condition is false

type IfExpr struct {
	Condition   Expression
	Consequence Expression
	Alternative Expression
}

func (i *IfExpr) String() string {
	var out strings.Builder
	out.WriteString("(if ")
	out.WriteString(i.Condition.String())
	out.WriteString(" ")
	out.WriteString(i.Consequence.String())
	if i.Alternative != nil {
		out.WriteString(" ")
		out.WriteString(i.Alternative.String())
	}
	out.WriteString(")")
	return out.String()
}

func (i *IfExpr) expressionNode() {}

// CondExpr represents a cond special form
// (cond (test1 expr1) (test2 expr2) ... (else exprN)) evaluates each test until one is true
// according to R4RS specification

type CondExpr struct {
	Clauses []CondClause
}

// CondClause represents a single clause in a cond expression

type CondClause struct {
	Test   Expression
	Result Expression
	IsElse bool
}

func (c *CondExpr) String() string {
	var out strings.Builder
	out.WriteString("(cond")
	for _, clause := range c.Clauses {
		out.WriteString(" ")
		if clause.IsElse {
			out.WriteString("(else ")
			out.WriteString(clause.Result.String())
			out.WriteString(")")
		} else {
			out.WriteString("(")
			out.WriteString(clause.Test.String())
			out.WriteString(" ")
			out.WriteString(clause.Result.String())
			out.WriteString(")")
		}
	}
	out.WriteString(")")
	return out.String()
}

func (c *CondExpr) expressionNode() {}

// ComparisonExpr represents a comparison expression
// (> expr1 expr2), (< expr1 expr2), etc.

type ComparisonExpr struct {
	Operator string
	Operands []Expression
}

func (ce *ComparisonExpr) String() string {
	var out strings.Builder
	out.WriteString("(")
	out.WriteString(ce.Operator)
	for _, operand := range ce.Operands {
		out.WriteString(" ")
		out.WriteString(operand.String())
	}
	out.WriteString(")")
	return out.String()
}

func (ce *ComparisonExpr) expressionNode() {}

// LogicalExpr represents a logical expression
// (and expr1 expr2 ...), (or expr1 expr2 ...), (not expr)

type LogicalExpr struct {
	Operator string
	Operands []Expression
}

func (le *LogicalExpr) String() string {
	var out strings.Builder
	out.WriteString("(")
	out.WriteString(le.Operator)
	for _, operand := range le.Operands {
		out.WriteString(" ")
		out.WriteString(operand.String())
	}
	out.WriteString(")")
	return out.String()
}

func (le *LogicalExpr) expressionNode() {}

// QuoteExpr represents a quote special form
// (quote expr) or 'expr

type QuoteExpr struct {
	Expression Expression
}

func (q *QuoteExpr) String() string {
	return fmt.Sprintf("(quote %s)", q.Expression.String())
}

func (q *QuoteExpr) expressionNode() {}
